use std::ffi::c_void;
use std::ptr::{self, null_mut};
use std::sync::atomic::{AtomicPtr, Ordering};
use std::time::Duration;
use dinvk::{GetModuleHandle, GetProcAddress};
use dinvk::data::CONTEXT_DEBUG_REGISTERS_AMD64;
use windows_sys::Win32::Foundation::{NTSTATUS, STATUS_SINGLE_STEP};
use windows_sys::Win32::System::Diagnostics::Debug::{CONTEXT, EXCEPTION_RECORD, RtlCaptureContext};
use windows_sys::Win32::System::Memory::{PAGE_READONLY, PAGE_READWRITE, VirtualProtect};

pub type BOOLEAN = u8;

use crate::dprintln;
use crate::return_ssn   ::{return_functions_ssn, return_syscall_instruction_address};
use crate::wow64_callback::return_wow64_function_pointer;

type NtContinueFn = unsafe extern "system" fn(*mut CONTEXT, BOOLEAN) -> NTSTATUS;
static NT_CONTINUE_PTR: AtomicPtr<c_void> = AtomicPtr::new(ptr::null_mut());


fn get_nt_continue() -> Option<NtContinueFn> {
    let ptr = NT_CONTINUE_PTR.load(Ordering::Relaxed);
    if !ptr.is_null() {
        return Some(unsafe { std::mem::transmute(ptr) });
    }

    let ntdll = GetModuleHandle("NTDLL.DLL", None);

    if ntdll.is_null() {
        dprintln!("[+] Failed to get Module of Ntdll");
        return None;
    }

    let proc = GetProcAddress(ntdll, "NtContinue", None);

    if proc.is_null() {
        dprintln!("[+] Unable to find Process Address");
        return None;
    }

    NT_CONTINUE_PTR.store(proc as *mut c_void, Ordering::Relaxed);

    dprintln!("[+] Loaded NtContinue at {:p}", proc);

    Some(unsafe { std::mem::transmute(proc) })
}

pub extern "system" fn exception_handler(exception_record: *mut EXCEPTION_RECORD, context_record: *mut CONTEXT) -> *mut c_void {

    let Some(nt_continue) = get_nt_continue() else {
        println!("[-] NtContinue Failed at Some");
        return ptr::null_mut();
    };

    unsafe {
        dprintln!("[*] Handler called with ExceptionCode: 0x{:x}, Rip: 0x{:x}, Dr0: 0x{:x}", (*exception_record).ExceptionCode, (*context_record).Rip, (*context_record).Dr0);
        
        if (*exception_record).ExceptionCode == STATUS_SINGLE_STEP {
            if (*context_record).Rip == (*context_record).Dr0 {
                dprintln!("[*] Condition matched, clearing Dr0, setting Rax to 0x{:x}, Rip to 0x{:x}", return_functions_ssn((*context_record).Rip), return_syscall_instruction_address((*context_record).Rip));
                (*context_record).Dr0 = 0;
                (*context_record).Rax = return_functions_ssn((*context_record).Rip);
                (*context_record).Rip = return_syscall_instruction_address((*context_record).Rip);
                let _status = nt_continue(context_record, 0);
                dprintln!("[*] NtContinue in handler returned 0x{:x} (should not reach here if successful)", _status);  // This should not print if successful
            } else {
                dprintln!("[*] Rip != Dr0, not handling");
            }
        } else {
            dprintln!("[*] Not STATUS_SINGLE_STEP, not handling");
        }
    }
    ptr::null_mut()
}

// For compiler optimization. use this = ) 

#[inline(never)]
pub fn set_hardware_breakpoint(nt_function_address: u64) {
    let Some(nt_continue) = get_nt_continue() else {
        println!("[ - ] Failed to load NtContinue");
        return;
    };
    dprintln!("[*] NtContinue function ready at {:p}", nt_continue as *const ());

    let inline_hook_address: u64 = nt_function_address + 3;
    dprintln!("[*] Calculated inline_hook_address: 0x{:x}", inline_hook_address);

    unsafe {
        let opcode = *(inline_hook_address as *const u8);
        dprintln!("[*] Opcode at inline_hook_address: 0x{:x}", opcode);
        if opcode != 0xE9 {
            println!("[ * ] Instruction at address 0x{:x} is not hooked", inline_hook_address);
            return;
        }
    }

    println!("[ * ] Setting hardware breakpoint at address: 0x{:x}", inline_hook_address);
    
    std::thread::sleep(Duration::from_secs(1));
    
    dprintln!("[*] Zeroing context");
    let mut context: CONTEXT = unsafe { std::mem::zeroed() };

    dprintln!("[*] Before RtlCaptureContext");
    unsafe { RtlCaptureContext(&mut context) };
    dprintln!("[*] After RtlCaptureContext, captured Rip: 0x{:x}", context.Rip);

    dprintln!("[*] Setting ContextFlags to 0x{:x}", CONTEXT_DEBUG_REGISTERS_AMD64);
    context.ContextFlags = CONTEXT_DEBUG_REGISTERS_AMD64;

    dprintln!("[*] Setting Dr0 to 0x{:x}", inline_hook_address);
    context.Dr0 = inline_hook_address;

    dprintln!("[*] Setting Dr7 to 0x00000001");
    context.Dr7 = 0x00000001;

    dprintln!("[*] Finish Setting Up Contexts");

    dprintln!("[*] Calling NtContinue");
    unsafe { nt_continue(&mut context, 0) };
}


pub fn hook_exception_dispatcher(module_base: *const u8) -> *mut c_void {
    let wow64_prepare_for_exception = return_wow64_function_pointer(module_base);
    if wow64_prepare_for_exception.is_null() {
        println!("[ - ] Failed to get Wow64PrepareForException function pointer.");
        return ptr::null_mut();
    }

    println!("[ * ] Wow64PrepareForException address: {:p}", wow64_prepare_for_exception);

    let mut old_protect: u32 = 0;

    let success = unsafe {
        VirtualProtect(
            wow64_prepare_for_exception as *const c_void,
            std::mem::size_of::<*mut c_void>(),
            PAGE_READWRITE,
            &mut old_protect,
        )
    };

    if success == 0 {
        return ptr::null_mut();
    }

    unsafe {
        *(wow64_prepare_for_exception as *mut usize) = exception_handler as usize;
    }

    let success = unsafe {
        VirtualProtect(
            wow64_prepare_for_exception as *const c_void,
            std::mem::size_of::<*mut c_void>(),
            PAGE_READONLY,
            &mut old_protect,
        )
    };

    if success == 0 {
        return null_mut();
    }

    println!("[ * ] pointer to ExceptionHandler written to: {:p}", wow64_prepare_for_exception);

    wow64_prepare_for_exception as _

}

pub fn unhook_exception_dispatcher(wow64_prepare_for_exception: *mut c_void) -> bool {
    if wow64_prepare_for_exception.is_null() {
        return false;
    }

    let mut old_protect: u32 = 0;
    let success = unsafe {
        VirtualProtect(
            wow64_prepare_for_exception,
            std::mem::size_of::<*mut c_void>(),
            PAGE_READWRITE,
            &mut old_protect,
        )
    };
    if success == 0 {
        return false;
    }

    unsafe {
        *(wow64_prepare_for_exception as *mut *mut c_void) = ptr::null_mut();
    }

    let success = unsafe {
        VirtualProtect(
            wow64_prepare_for_exception,
            std::mem::size_of::<*mut c_void>(),
            PAGE_READONLY,
            &mut old_protect,
        )
    };
    if success == 0 {
        return false;
    }

    true
}